"""
optimal_knn_webapp_pinecone.py - 使用 Pinecone 的手写数字识别 Web 应用

修改自 IPR 问题 2: 使用最佳 KNN 模型识别手写数字的 APP 的 optimal_knn_webapp.py 文件
使其能够实现第一次推理速度较慢，但是第二次推理速度较快的效果

要求：
- 修改原有的 optimal_knn_webapp.py 文件
- 使用 Pinecone 云端向量数据库进行推理
- 第一次推理会较慢（需要连接 Pinecone）
- 第二次及之后的推理会更快（已建立连接）
"""

from __future__ import annotations

import io
import time
from pathlib import Path
from typing import Tuple, Any, Optional

import numpy as np
from PIL import Image, ImageOps, ImageFilter

import gradio as gr
try:
    import cv2  # 可选：用于更快的 Otsu 与膨胀
except Exception:
    cv2 = None

try:
    from pinecone import Pinecone
    PINECONE_AVAILABLE = True
except ImportError:
    PINECONE_AVAILABLE = False
    print("警告: pinecone 未安装。请运行: pip install pinecone")

# Pinecone 配置
PINECONE_API_KEY = "pcsk_6xFLgp_DBvvvh8Yzqz3qD3rH4KykRKBakV3Br9BtuQjUnccSNNuJ59BoaGzvLGtVRnocd9"
INDEX_NAME = "mnist-index"

# 全局变量：Pinecone 连接（第一次推理时初始化，之后复用）
_PINECONE_CLIENT = None
_PINECONE_INDEX = None
_FIRST_INFERENCE = True


def init_pinecone():
    """
    初始化 Pinecone 连接（仅在第一次推理时调用）
    """
    global _PINECONE_CLIENT, _PINECONE_INDEX, _FIRST_INFERENCE
    
    if not PINECONE_AVAILABLE:
        raise ImportError("请先安装 pinecone: pip install pinecone")
    
    if _PINECONE_CLIENT is None:
        print("正在初始化 Pinecone 连接...")
        _PINECONE_CLIENT = Pinecone(api_key=PINECONE_API_KEY)
        _PINECONE_INDEX = _PINECONE_CLIENT.Index(INDEX_NAME)
        print("Pinecone 连接成功！")
        _FIRST_INFERENCE = False


def _to_pil(input_obj: Any) -> Image.Image:
    """将 Sketchpad 可能返回的多种类型(dict/ndarray/PIL)统一转为 PIL.Image.

    兼容格式：
    - PIL.Image
    - numpy.ndarray (H,W) 或 (H,W,3|4)
    - dict：常见键 image/composite/background/layers（取首个可解析的图层）
    """
    if input_obj is None:
        raise ValueError("未检测到图像")

    # 已是 PIL
    if isinstance(input_obj, Image.Image):
        pil = input_obj
    # numpy 数组
    elif isinstance(input_obj, np.ndarray):
        pil = Image.fromarray(input_obj)
    # gradio 某些版本的 Sketchpad/ImageEditor 可能返回 dict
    elif isinstance(input_obj, dict):
        candidate = None
        for key in ("image", "composite", "background"):
            if key in input_obj and input_obj[key] is not None:
                candidate = input_obj[key]
                break
        # 尝试 layers（若存在），取最后一层
        if candidate is None and isinstance(input_obj.get("layers"), list) and input_obj["layers"]:
            last = input_obj["layers"][-1]
            if isinstance(last, dict):
                candidate = last.get("image") or last.get("composite")
            else:
                candidate = last

        if candidate is None:
            raise ValueError("无法从输入字典解析图像内容")

        if isinstance(candidate, Image.Image):
            pil = candidate
        elif isinstance(candidate, np.ndarray):
            pil = Image.fromarray(candidate)
        else:
            # 有些实现可能是 base64/bytes，这里先简单兜底
            raise ValueError("不支持的图像数据类型: " + type(candidate).__name__)
    else:
        raise ValueError("不支持的输入类型: " + type(input_obj).__name__)

    # 去除 alpha：合成到白底，避免反相后出现暗底
    if pil.mode in ("RGBA", "LA"):
        bg = Image.new("RGBA", pil.size, (255, 255, 255, 255))
        pil = Image.alpha_composite(bg, pil.convert("RGBA")).convert("RGB")

    return pil


def _otsu_threshold(arr: np.ndarray) -> int:
    """简易 Otsu 自动阈值计算，返回 0..255 的整型阈值。"""
    # 优先使用 OpenCV 的 Otsu（C 实现更快）
    if cv2 is not None:
        try:
            _arr = arr.astype(np.uint8, copy=False)
            _, th = cv2.threshold(_arr, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            return int(_)
        except Exception:
            pass
    hist, _ = np.histogram(arr, bins=256, range=(0, 255))
    total = arr.size
    sum_total = np.dot(np.arange(256), hist)
    sumB = 0.0
    wB = 0.0
    max_var = -1.0
    thresh = 0
    for t in range(256):
        wB += hist[t]
        if wB == 0:
            continue
        wF = total - wB
        if wF == 0:
            break
        sumB += t * hist[t]
        mB = sumB / wB
        mF = (sum_total - sumB) / wF
        varBetween = wB * wF * (mB - mF) ** 2
        if varBetween > max_var:
            max_var = varBetween
            thresh = t
    return int(thresh)


def _preprocess_to_digits(
    img: Image.Image,
    invert: bool = True,
    thresh: int = 10,
    padding: int = 4,
    blur_radius: float = 1.0,
    dilate: int = 0,
    auto_thresh: bool = False,
) -> np.ndarray:
    """
    将任意输入图像预处理为 sklearn digits (8x8, 值域 0..16) 的扁平 1x64 向量。

    步骤：转灰度 -> 反相(使数字为亮、背景暗) -> 缩放到 8x8 -> 归一化到 0..16 -> flatten
    """
    # 1. 转灰度
    gray = ImageOps.grayscale(img)
    arr = np.array(gray, dtype=np.float32)

    # 2. 反相（可选）
    if invert:
        arr = 255.0 - arr

    # 3. 自动阈值（可选）
    if auto_thresh:
        thresh = _otsu_threshold(arr.astype(np.uint8))

    # 4. 二值化 + 找前景
    binary = (arr > thresh).astype(np.uint8) * 255

    # 5. 膨胀（可选）
    if dilate > 0:
        if cv2 is not None:
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
            for _ in range(dilate):
                binary = cv2.dilate(binary, kernel)
        else:
            # 简易膨胀：最大池化
            from scipy.ndimage import maximum_filter
            for _ in range(dilate):
                binary = maximum_filter(binary, size=3)

    # 6. 找外接框
    coords = np.column_stack(np.where(binary > 0))
    if coords.size == 0:
        # 空白图像，返回全零向量
        return np.zeros(64, dtype=np.float32)

    y_min, x_min = coords.min(axis=0)
    y_max, x_max = coords.max(axis=0)

    # 7. 裁剪 + padding
    y_min = max(0, y_min - padding)
    x_min = max(0, x_min - padding)
    y_max = min(binary.shape[0], y_max + padding + 1)
    x_max = min(binary.shape[1], x_max + padding + 1)

    cropped = binary[y_min:y_max, x_min:x_max]

    # 8. 缩放到 8x8
    pil_crop = Image.fromarray(cropped, mode="L")
    if blur_radius > 0:
        pil_crop = pil_crop.filter(ImageFilter.GaussianBlur(radius=blur_radius))
    pil_8x8 = pil_crop.resize((8, 8), Image.LANCZOS)

    # 9. 归一化到 0..16
    arr_8x8 = np.array(pil_8x8, dtype=np.float32)
    arr_8x8 = arr_8x8 / 255.0 * 16.0

    # 10. flatten
    return arr_8x8.flatten()


def predict_digit_with_pinecone(
    input_img: Any,
    invert: bool = True,
    thresh: int = 10,
    padding: int = 4,
    blur_radius: float = 1.0,
    dilate: int = 0,
    preview_size: int = 128,
    show_grid: bool = False,
    auto_thresh: bool = True,
    fast_mode: bool = False,
    k: int = 11,
) -> Tuple[int, str, Optional[Image.Image]]:
    """
    使用 Pinecone 进行预测
    
    第一次调用会较慢（需要初始化 Pinecone 连接）
    之后的调用会更快（复用已建立的连接）
    """
    global _FIRST_INFERENCE
    
    try:
        # 第一次推理：初始化 Pinecone
        start_time = time.time()
        
        if _PINECONE_INDEX is None:
            init_pinecone()
            init_info = "（首次推理，正在连接 Pinecone...）"
        else:
            init_info = "（使用已建立的 Pinecone 连接）"
        
        if input_img is None:
            return -1, "未检测到图像，请在画板中写一个数字后再提交。", None

        # 预处理图像
        pil_img = _to_pil(input_img)
        if fast_mode:
            features = _preprocess_to_digits(
                pil_img,
                invert=invert,
                thresh=thresh,
                padding=padding,
                blur_radius=0.0,
                dilate=0,
                auto_thresh=False,
            )
        else:
            features = _preprocess_to_digits(
                pil_img,
                invert=invert,
                thresh=thresh,
                padding=padding,
                blur_radius=blur_radius,
                dilate=dilate,
                auto_thresh=auto_thresh,
            )
        
        # 使用 Pinecone 查询
        query_start = time.time()
        results = _PINECONE_INDEX.query(
            vector=features.tolist(),
            top_k=k,
            include_metadata=True
        )
        query_time = time.time() - query_start
        
        # 获取最近邻的标签并投票
        labels = [match['metadata']['label'] for match in results['matches']]
        predicted_label = int(np.bincount(labels).argmax())
        
        # 计算总时间
        total_time = time.time() - start_time
        
        # 生成说明信息
        info = f"预测: {predicted_label} (K={k}, Pinecone 查询)\n"
        info += f"{init_info}\n"
        info += f"总耗时: {total_time:.3f}秒 (查询: {query_time:.3f}秒)"
        
        # 生成 8x8 预处理可视化图
        img8 = (features.reshape(8, 8) / 16.0 * 255.0).astype(np.uint8)
        pv = 128 if fast_mode else max(32, int(preview_size))
        viz = Image.fromarray(img8, mode="L").resize((pv,) * 2, Image.NEAREST)
        
        if show_grid:
            # 叠加网格，帮助观察每个像素格
            viz = viz.convert("RGB")
            import PIL.ImageDraw as ImageDraw
            draw = ImageDraw.Draw(viz)
            step = viz.size[0] // 8
            for i in range(1, 8):
                draw.line([(i * step, 0), (i * step, viz.size[1])], fill=(200, 200, 200), width=1)
                draw.line([(0, i * step), (viz.size[0], i * step)], fill=(200, 200, 200), width=1)
        
        return predicted_label, info, viz
        
    except Exception as e:
        return -1, f"发生错误：{e}", None


# 创建 Gradio 接口
with gr.Blocks(title="KNN 手写数字识别 (Pinecone 版本)") as demo:
    gr.Markdown("""
    # KNN 手写数字识别 (Pinecone 版本)
    
    在画板中写下一个数字并提交，模型将使用 Pinecone 云端向量数据库给出预测结果。
    
    **特点**：
    - 第一次推理较慢（需要连接 Pinecone）
    - 第二次及之后的推理更快（复用连接）
    """)

    with gr.Row():
        sketch = gr.Sketchpad(
            label="画板(建议写粗一点，填满数字)",
            type="numpy",
            height=280,
            width=280,
        )
        with gr.Column():
            gr.Markdown("### 📝 提示：写粗一点，数字要饱满！")
            invert_chk = gr.Checkbox(value=True, label="反相(白字黑底) - 必须开启")
            with gr.Row():
                thresh_sld = gr.Slider(0, 255, value=50, step=1, label="阈值(找前景) - 建议50-80")
                pad_sld = gr.Slider(0, 32, value=8, step=1, label="外接框边距 - 建议6-10")
            with gr.Row():
                blur_sld = gr.Slider(0, 3, value=1.0, step=0.5, label="模糊半径 - 建议1.0")
                dilate_sld = gr.Slider(0, 5, value=2, step=1, label="膨胀(加粗) - 建议2-3")
            with gr.Row():
                prev_sld = gr.Slider(64, 256, value=192, step=32, label="预览尺寸")
                grid_chk = gr.Checkbox(value=True, label="显示网格 - 建议开启")
            with gr.Row():
                k_sld = gr.Slider(1, 20, value=11, step=1, label="K 值 - 建议9-13")
            auto_thresh_chk = gr.Checkbox(value=True, label="自动阈值(Otsu) - 建议开启")
            fast_chk = gr.Checkbox(value=False, label="极速模式")
            
            gr.Markdown("""
            **🎯 提高准确率的技巧：**
            1. 写粗一点，让数字更饱满
            2. 数字要居中，占据大部分画板
            3. 如果识别错误，尝试调整阈值(50-80)
            4. 增加膨胀值(2-3)可以让数字更粗
            """)
            
            btn = gr.Button("🚀 提交 (使用 Pinecone)", variant="primary")
            out_label = gr.Number(label="预测数字", precision=0)
            out_text = gr.Textbox(label="说明", lines=3)
            out_preview = gr.Image(label="预处理8x8视图 (查看是否清晰)", type="pil")

            # 清除按钮
            try:
                gr.ClearButton(
                    components=[sketch, out_preview],
                    value="清除",
                )
            except Exception:
                clear_btn = gr.Button("清除")
                def _clear_only_canvas_preview():
                    return None, None
                clear_btn.click(
                    _clear_only_canvas_preview,
                    inputs=None,
                    outputs=[sketch, out_preview],
                )

    btn.click(
        predict_digit_with_pinecone,
        inputs=[sketch, invert_chk, thresh_sld, pad_sld, blur_sld, dilate_sld, 
                prev_sld, grid_chk, auto_thresh_chk, fast_chk, k_sld],
        outputs=[out_label, out_text, out_preview],
    )


if __name__ == "__main__":
    print("="*60)
    print("KNN 手写数字识别 Web 应用 (Pinecone 版本)")
    print("="*60)
    print("注意：")
    print("1. 第一次推理会较慢（需要连接 Pinecone）")
    print("2. 之后的推理会更快（复用已建立的连接）")
    print("3. 请确保已运行 pinecone_train.py 创建索引")
    print("="*60)
    demo.launch()
